package z.nat;

import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.http.DefaultHttpHeaders;
import io.netty.handler.codec.http.HttpClientCodec;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.websocketx.*;
import io.netty.handler.timeout.IdleStateEvent;
import io.netty.handler.timeout.IdleStateHandler;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URI;
import java.util.concurrent.TimeUnit;

/**
 * 隧道客户端，部署在内网环境中
 * 1. 连接到公网服务器的WebSocket服务
 * 2. 接收服务端的请求并转发到本地VNC服务
 */
public class WebSocketClient extends SimpleChannelInboundHandler<WebSocketFrame> {
    private static final Logger logger = LoggerFactory.getLogger(WebSocketClient.class);

    private static String serverHost = "xajysj.com", vncHost = "127.0.0.1";
    private static int serverPort = 40080, vncPort = 2222, proxyPort = 43322;

    public static void main(String[] args) throws Exception {
        // 解析命令行参数
        if (args.length > 4) {
            serverHost = args[0];
            serverPort = Integer.parseInt(args[1]);
            vncHost = args[2];
            vncPort = Integer.parseInt(args[3]);
            proxyPort = Integer.parseInt(args[4]);
        }

        logger.info("启动隧道客户端，连接到服务端: {}:{}, 本地VNC: {}:{}",
                serverHost, serverPort, vncHost, vncPort);

        new WebSocketClient();
    }

    public static EventLoopGroup eventLoopGroup;

    private WebSocketClient() throws Exception {
        eventLoopGroup = new NioEventLoopGroup(1);

        // 构建WebSocket URI
        URI uri = new URI("wss://" + serverHost + ":" + serverPort + "/tunnel");

        // 创建WebSocket客户端
        Bootstrap bootstrap = new Bootstrap().group(eventLoopGroup)
                .channel(NioSocketChannel.class)
                .option(ChannelOption.SO_KEEPALIVE, true)
                .handler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) {
                        ChannelPipeline pipeline = ch.pipeline();
                        pipeline.addLast(
                                new HttpClientCodec(),
                                new HttpObjectAggregator(65536),
                                new WebSocketClientProtocolHandler(
                                        WebSocketClientHandshakerFactory.newHandshaker(
                                                uri, WebSocketVersion.V13, null, false, new DefaultHttpHeaders()
                                        )
                                ),
                                new IdleStateHandler(0, 30, 0, TimeUnit.SECONDS),
                                WebSocketClient.this
                        );
                    }
                });

        // 连接到服务端
        Channel channel = bootstrap.connect(serverHost, serverPort).sync().channel();
        logger.info("连接到服务端: {}:{}", serverHost, serverPort);

        // 等待连接关闭
        channel.closeFuture().sync();
    }


    public static void stop() {
        if (eventLoopGroup == null) return;
        eventLoopGroup.shutdownGracefully();
        eventLoopGroup = null;
    }


    private VncProxyClient vncProxyClient;

    // 不在channelActive中发送注册消息，因为此时WebSocket握手尚未完成
    @Override
    public void channelActive(ChannelHandlerContext ctx) {
        logger.info("TCP连接已建立，等待WebSocket握手完成");
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, WebSocketFrame frame) {
        if (frame instanceof TextWebSocketFrame) {
            handleTextFrame(ctx, (TextWebSocketFrame) frame);
        } else if (frame instanceof BinaryWebSocketFrame) {
            handleBinaryFrame(ctx, (BinaryWebSocketFrame) frame);
        } else if (frame instanceof PingWebSocketFrame) {
            ctx.writeAndFlush(new PongWebSocketFrame(frame.content().retain()));
        } else if (frame instanceof CloseWebSocketFrame) {
            ctx.close();
        }
    }


    private void closeVncConnection() {
        if (vncProxyClient == null) return;
        vncProxyClient.closeVncConnection();
        vncProxyClient = null;
    }

    /**
     * 处理文本帧，用于控制命令
     */
    private void handleTextFrame(ChannelHandlerContext ctx, TextWebSocketFrame frame) {
        String text = frame.text();
        logger.debug("收到WebSocket文本消息: {}", text);

        switch (text) {
            // 处理心跳消息
            case "PING":
                ctx.writeAndFlush(new TextWebSocketFrame("PONG"));
                return;
            case "CONNECT":
                closeVncConnection();
                vncProxyClient = new VncProxyClient(ctx, vncHost, vncPort);
                return;
            case "CLOSE":
                closeVncConnection();
        }
    }

    /**
     * 处理二进制帧，用于VNC数据传输
     */
    private void handleBinaryFrame(ChannelHandlerContext ctx, BinaryWebSocketFrame frame) {
        ByteBuf content = frame.content();
        // 转发数据到VNC连接
        if (vncProxyClient == null) return;
        Channel channel1 = vncProxyClient.getVncChannel();
        if (channel1 == null || !channel1.isActive()) return;
        channel1.writeAndFlush(content.retain());
    }

    private boolean reconnecting;


    @Override
    public void channelInactive(ChannelHandlerContext ctx) {
        logger.info("WebSocket连接断开，关闭VNC连接");
        closeVncConnection();

        // 如果不是正在重连，则尝试重连
        if (!reconnecting) {
            reconnecting = true;
            ctx.channel().eventLoop().schedule(() -> {
                logger.info("尝试重新连接到服务端...");
                try {
                    new WebSocketClient();
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
                reconnecting = false;
            }, 5, TimeUnit.SECONDS);
        }
    }

    @Override
    public void userEventTriggered(ChannelHandlerContext ctx, Object evt) {
        if (evt instanceof IdleStateEvent) {
            // 空闲超时，发送心跳检测
            ctx.writeAndFlush(new TextWebSocketFrame("PING"));
            logger.debug("发送心跳检测");
        } else if (evt == WebSocketClientProtocolHandler.ClientHandshakeStateEvent.HANDSHAKE_COMPLETE) {
            // WebSocket握手完成后，发送注册消息
            logger.info("WebSocket握手完成，正在注册客户端");
            ctx.writeAndFlush(new TextWebSocketFrame("REGISTER:" + proxyPort));
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        logger.error("WebSocket客户端异常", cause);
        ctx.close();
    }

}